package org.checkerframework.framework.source;

import com.sun.source.tree.ClassTree;
import com.sun.source.tree.CompilationUnitTree;
import com.sun.source.tree.MethodTree;
import com.sun.source.tree.Tree;
import com.sun.source.tree.VariableTree;
import com.sun.source.util.TreePath;
import com.sun.source.util.TreePathScanner;
import com.sun.source.util.Trees;
import java.util.ArrayList;
import java.util.List;
import javax.annotation.processing.ProcessingEnvironment;
import javax.lang.model.element.Element;
import javax.lang.model.util.Elements;
import javax.lang.model.util.Types;
import org.checkerframework.javacutil.BugInCF;
import org.checkerframework.javacutil.TreeUtils;

/**
 * An AST visitor that provides a variety of compiler utilities and interfaces to facilitate
 * type-checking.
 */
public abstract class SourceVisitor<R, P> extends TreePathScanner<R, P> {

  /** The {@link Trees} instance to use for scanning. */
  protected final Trees trees;

  /** The {@link Elements} helper to use when scanning. */
  protected final Elements elements;

  /** The {@link Types} helper to use when scanning. */
  protected final Types types;

  /**
   * The root of the AST that this {@link SourceVisitor} will scan.
   *
   * <p>Is set by {@link #setRoot}.
   */
  protected CompilationUnitTree root;

  /** The trees that are annotated with {@code @SuppressWarnings}. */
  public final List<Tree> treesWithSuppressWarnings;

  /** True if a warning should be issued for unneeded warning suppressions. */
  private final boolean warnUnneededSuppressions;

  /**
   * Creates a {@link SourceVisitor} to use for scanning a source tree.
   *
   * @param checker the checker to invoke on the input source tree
   */
  protected SourceVisitor(SourceChecker checker) {
    // Use the checker's processing environment to get the helpers we need.
    ProcessingEnvironment env = checker.getProcessingEnvironment();

    this.trees = Trees.instance(env);
    this.elements = env.getElementUtils();
    this.types = env.getTypeUtils();

    this.treesWithSuppressWarnings = new ArrayList<>();
    this.warnUnneededSuppressions = checker.hasOption("warnUnneededSuppressions");
  }

  /**
   * Set the CompilationUnitTree to be used during any visits. For any later calls of {@code
   * com.sun.source.util.TreePathScanner.scan(TreePath, P)}, the CompilationUnitTree of the TreePath
   * has to be equal to {@code root}.
   *
   * @param newRoot the new compilation unit
   */
  public void setRoot(CompilationUnitTree newRoot) {
    this.root = newRoot;
  }

  /**
   * Store the last Tree visited by the SourceVisitor. This is necessary because the finally blocks
   * in {@link com.sun.source.util.TreePathScanner#scan(TreePath, Object)} and {@link
   * com.sun.source.util.TreePathScanner#scan(Tree, Object)} set the visited Path to null. This
   * field is used to report a rough location for the error in {@link
   * org.checkerframework.framework.source.SourceChecker#logBugInCF(BugInCF)}.
   */
  /*package-private*/ Tree lastVisited;

  /** Entry point for a type processor: the TreePath leaf is a top-level type tree within root. */
  public void visit(TreePath path) {
    lastVisited = path.getLeaf();
    this.scan(path, null);
  }

  @Override
  public R scan(Tree tree, P p) {
    lastVisited = tree;
    return super.scan(tree, p);
  }

  @Override
  public R visitClass(ClassTree classTree, P p) {
    storeSuppressWarningsAnno(classTree);
    return super.visitClass(classTree, p);
  }

  @Override
  public R visitVariable(VariableTree variableTree, P p) {
    storeSuppressWarningsAnno(variableTree);
    return super.visitVariable(variableTree, p);
  }

  @Override
  public R visitMethod(MethodTree tree, P p) {
    storeSuppressWarningsAnno(tree);
    return super.visitMethod(tree, p);
  }

  /**
   * If {@code tree} has a {@code @SuppressWarnings} add it to treesWithSuppressWarnings.
   *
   * @param tree a declaration on which a {@code @SuppressWarnings} annotation may be placed
   */
  private void storeSuppressWarningsAnno(Tree tree) {
    if (!warnUnneededSuppressions) {
      return;
    }
    Element elt = TreeUtils.elementFromTree(tree);
    if (elt.getAnnotation(SuppressWarnings.class) != null) {
      treesWithSuppressWarnings.add(tree);
    }
  }
}
